# Real-world Scenario Environment Visualization

import argparse
import random

from metadrive.policy.replay_policy import ReplayEgoCarPolicy
from metadrive.constants import HELP_MESSAGE
from metadrive.engine.asset_loader import AssetLoader
from metadrive.envs import ScenarioEnv

class DemoEnv(ScenarioEnv):
    def reset(self, seed=None):
        if self.engine is not None:
            seeds = [i for i in range(self.config["num_scenarios"])]
            seeds.remove(self.current_seed)
            seed = random.choice(seeds)
        return super(DemoEnv, self).reset(seed=seed)
    
extra_args = dict(film_size=(1200, 1200))
asset_path = AssetLoader.asset_path

try:
    env = DemoEnv(
        {
            "manual_control": False,
            "reactive_traffic": False,
            "use_render": False,
            "agent_policy": ReplayEgoCarPolicy,
            "data_directory": AssetLoader.file_path(asset_path, "waymo", unix_style=False),
            "num_scenarios": 3
        }
    )

    o, _ = env.reset()
    frames = []
    for i in range(1, 100000):
        o, r, tm, tc, info =env.step([1.0, 0.])
        frame = env.render(
            mode="top_down",
            **extra_args
        )
        frames.append(frame)
        if tm or tc:
            env.reset()
            break
        
except Exception as e:
    raise e
finally:
    env.close()

# render image
print(f"Generate gif ...")
import pygame
import numpy as np

from PIL import Image
import os

HOME = os.path.dirname(os.path.abspath(__file__))

images = [frame for frame in frames]
images = [Image.fromarray(frame) for frame in images]
images[0].save(f"{HOME}/real_world.gif", save_all=True, append_images=images[1:], loop=0)
